
import numpy as np
import os
import PIL.Image

from .cut_alg import get_patches


def split_array(src_hwc_arr, dst_hw, overlap_hw):
    src_hw = src_hwc_arr.shape[0:2]
    for h_offset, w_offset, h_end, w_end in get_patches(src_hw, dst_hw, overlap_hw):
        yield src_hwc_arr[h_offset:h_end, w_offset:w_end], (h_offset, w_offset)


def main(args):
    PIL.Image.MAX_IMAGE_PIXELS = None

    src_basename = os.path.basename(args.src)

    src = PIL.Image.open(args.src)
    dst_hw = list(map(int, args.dst_hw.split(',')))
    overlap_factor = args.overlap_factor
    overlap_hw = tuple(int(x * overlap_factor) for x in dst_hw)

    if args.dryrun:
        src_hw = src.size
        for h_offset, w_offset, h_end, w_end in get_patches(src_hw, dst_hw, overlap_hw):
            print(h_offset, w_offset, h_end, w_end)

        return

    if args.dst_dir is None:
        args.dst_dir = os.path.dirname(args.src)
    os.makedirs(args.dst_dir, exist_ok=True)

    src_hwc_arr = np.array(src)
    for i, (h_, hw_offset) in enumerate(split_array(src_hwc_arr, dst_hw, overlap_hw)):
        key = os.path.join(args.dst_dir, src_basename)
        dst_filename = '{}__{}___{}.png'.format(key, hw_offset[0], hw_offset[1])
        dst = PIL.Image.fromarray(dst_hwc_arr)
        dst.save(dst_filename)


if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('src', type=str)
    parser.add_argument('--dst-dir', type=str)
    parser.add_argument('--dst-hw', type=str, default='640,640')
    parser.add_argument('--overlap-factor', type=float, default=0.25)
    parser.add_argument('--dryrun', action='store_true')
    args = parser.parse_args()
    main(args)
